package com.caoyanan.algorithm.question.zuoTraining.struct;

import java.util.HashMap;
import java.util.Map;
import java.util.Stack;

/**
 * @author: caoyanan
 * @time: 2021/5/20 5:46 下午
 */
public class UnionFind {

    private final Map<Integer, Integer> parent;
    private final Map<Integer, Integer> size;

    public UnionFind(int[] arr) {
        parent = new HashMap<>(arr.length);
        size = new HashMap<>(arr.length);
        for (int i = 0; i < arr.length; i++) {
            parent.put(arr[i], arr[i]);
            size.put(arr[i], 1);
        }
    }

    private int findParent(int i) {
        Stack<Integer> path = new Stack<>();
        while (i != parent.get(i)) {
            path.push(parent.get(i));
            i = parent.get(i);
        }
        while (!path.isEmpty()) {
            parent.put(path.pop(), i);
        }
        return i;
    }

    public boolean isSame(int i, int j) {
        return findParent(i) == findParent(j);
    }

    public void union(int i, int j) {
        int iParent = findParent(i);
        int jParent = findParent(j);
        if (iParent == jParent) {
            return;
        }
        int big = size.get(iParent) > size.get(jParent) ? iParent : jParent;
        int small = big == iParent ? jParent : iParent;
        parent.put(small, big);
        size.put(big, size.get(big) + size.get(small));
        size.remove(small);
    }

    public int size() {
        return size.size();
    }
}
